from tree_gui import Ui_Form
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QMessageBox, QFileDialog
from PyQt5 import QtWidgets, QtCore, QtGui
import matplotlib
matplotlib.use('qt5Agg')
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import matplotlib.pyplot as plt
import networkx as nx
from BST import *
from AVL import *

def create_graph(G, node, pos={}, x=0, y=0, layer=1):
    pos[node.data] = (x, y)
    if node.lchild:
        G.add_edge(node.data, node.lchild.data)
        l_x, l_y = x - 1 / 2 ** layer, y - 1
        l_layer = layer + 1
        create_graph(G, node.lchild, x=l_x, y=l_y, pos=pos, layer=l_layer)
    if node.rchild:
        G.add_edge(node.data, node.rchild.data)
        r_x, r_y = x + 1 / 2 ** layer, y - 1
        r_layer = layer + 1
        create_graph(G, node.rchild, x=r_x, y=r_y, pos=pos, layer=r_layer)
    return (G, pos)

class Demo(QWidget, Ui_Form):
    def __init__(self):
        super(Demo, self).__init__()
        self.setupUi(self)
        self.figure = plt.figure()
        self.canvas = FigureCanvas(self.figure)
        self.layout_init()
        self.tree = AVLTree([5, 4, 6, 1])
        self.button_init()

    def layout_init(self):
        self.h1_layout = QtWidgets.QHBoxLayout()
        self.groupBox_5.setLayout(self.h1_layout)
        self.h1_layout.addWidget(self.canvas)

    def button_init(self):
        self.button_show.clicked.connect(self.show_tree)
        self.add_node.clicked.connect(self.add)
        self.update_add.clicked.connect(self.update)
        self.delete_node.clicked.connect(self.delete)
        self.update_delete.clicked.connect(self.update)
        self.search_node.clicked.connect(self.search)

    def show_tree(self):
        temp = self.initia_node.toPlainText()
        temp = eval(temp)
        self.tree = AVLTree(temp)
        self._plot()

    def add(self):
        val = eval(self.add_node_input.text())
        self.tree.insert_no_rec(val)
        QtWidgets.QMessageBox.information(self, 'information', '节点添加成功！')

    def update(self):
        self._plot()

    def delete(self):
        val = eval(self.delete_node_input.text())
        self.tree.delete(val)
        QtWidgets.QMessageBox.information(self, 'information', '节点删除成功！')

    def search(self):
        val = eval(self.search_node_input.text())
        print(val)
        node = self.tree.search_no_rec(val)
        if not node:
            print(1)
            QtWidgets.QMessageBox.critical(self, 'warning', '未找到该节点！')
        else:
            print(2)
            QtWidgets.QMessageBox.information(self, 'information', '该节点的值为：{}'.format(node.data))





    def _plot(self):
        ax = self.figure.add_axes([0.1, 0.1, 0.8, 0.8])  # ax1 = fig.add_axes([left, bottom, width, height])
        ax.cla()
        graph = nx.DiGraph()
        graph, pos = create_graph(graph, self.tree.root)
        nx.draw_networkx(graph, pos, ax=ax, node_size=300)
        self.canvas.draw()



if __name__=='__main__':
    app = QApplication(sys.argv)
    demo = Demo()
    demo.show()
    sys.exit(app.exec_())
